"""
    PascalVOC格式数据集转换为YOLO格式
"""

import argparse
import os
import xml.etree.ElementTree as ET
from tqdm import tqdm


def xyxy2xywh(imsize: list, xyxy: list) -> list:
    """
    trans xyxy to xywh
    :param imsize: size of image
    :param xyxy: [xmin, ymin, xmax, ymax]
    :return: [x_center, y_center, w, h] relative
    """
    width, height = imsize[0], imsize[1]
    x_center = ((xyxy[0] + xyxy[2]) / 2) / width
    y_center = ((xyxy[1] + xyxy[3]) / 2) / height
    w = (xyxy[2] - xyxy[0]) / width
    h = (xyxy[3] - xyxy[1]) / height

    return [round(xywh, 6) for xywh in [x_center, y_center, w, h]]


def main(voc_path: str, save_path: str) -> None:
    """

    :param voc_path: voc数据集目录的路径
    :param save_path: 保存的路径
    :return: None
    """
    classes = []

    # filter file endswith ".xml"
    voc_files = list(filter(lambda voc_file: voc_file.endswith(".xml"), os.listdir(voc_path)))

    for voc_file in tqdm(voc_files):
        f_yolo = open(os.path.join(save_path, voc_file[:-4] + ".txt"), "w")

        tree = ET.parse(os.path.join(voc_path, voc_file))
        root = tree.getroot()

        size = root.find('size')
        width = int(size.find('width').text)
        height = int(size.find('height').text)

        objects = root.findall('object')
        for obj in objects:
            name = obj.find('name').text
            if name not in classes:
                classes.append(name)

            idx = classes.index(name)

            bnd = obj.find('bndbox')
            xmin = bnd.find('xmin').text
            ymin = bnd.find('ymin').text
            xmax = bnd.find('xmax').text
            ymax = bnd.find('ymax').text

            coord = [int(xy) for xy in [xmin, ymin, xmax, ymax]]
            xywh = xyxy2xywh([width, height], coord)

            f_yolo.write(str(idx) + ' ')
            f_yolo.write(str(xywh[0]) + ' ')
            f_yolo.write(str(xywh[1]) + ' ')
            f_yolo.write(str(xywh[2]) + ' ')
            f_yolo.write(str(xywh[3]) + '\n')

        f_yolo.close()

def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument("--vocPath", type=str, default="dataset/labels/voc", help="voc dataset path")
    parser.add_argument("--savePath", type=str, default="dataset/labels/yolo", help="yolo dataset save path")
    return parser.parse_args()


if __name__ == '__main__':
    opt = parse_opt()
    main(opt.vocPath, opt.savePath)
